import re
from neo4j import GraphDatabase
import nltk
from nltk.tokenize import sent_tokenize
import random
from fastapi import FastAPI, Depends, HTTPException, status, Body
from fastapi.middleware.cors import CORSMiddleware
from sqlalchemy.orm import Session
from datetime import timedelta, datetime
from typing import Dict, List, Any, Optional
from contextlib import asynccontextmanager
import models, schemas, auth
from database import engine, get_db
import logging
import os
import json
import uuid
from dotenv import load_dotenv
from pydantic import BaseModel

# Import our trust model component
from mtmkg_model import MTMKGTrustModel

# 加载环境变量
load_dotenv()

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 下载 NLTK 数据
nltk.download('punkt')

# Neo4j 连接配置
URI = os.getenv("NEO4J_URI", "bolt://localhost:7687")
USER = os.getenv("NEO4J_USER", "neo4j")
PASSWORD = os.getenv("NEO4J_PASSWORD", "password")

logger.info(f"Neo4j 连接 URI: {URI}")

# 初始化信任模型组件
mtmkg_model = MTMKGTrustModel(URI, USER, PASSWORD)

# 辅助函数：处理Neo4j数据类型
def serialize_neo4j_data(data: Any) -> Any:
    """将Neo4j数据类型转换为可JSON序列化的类型"""
    if hasattr(data, "__class__") and data.__class__.__module__ == "neo4j.time":
        # 处理Neo4j的时间类型
        if hasattr(data, "iso_format"):
            return data.iso_format()  # DateTime转为ISO格式字符串
        return str(data)  # 其他时间类型转为字符串
    elif isinstance(data, dict):
        # 递归处理字典
        return {k: serialize_neo4j_data(v) for k, v in data.items()}
    elif isinstance(data, list):
        # 递归处理列表
        return [serialize_neo4j_data(item) for item in data]
    return data

class TrustEvaluationInput(BaseModel):
    source_id: str
    target_id: str
    relation_type: str

class KnowledgeGraph:
    def __init__(self, uri, user, password):
        try:
            logger.info(f"正在连接到 Neo4j 数据库: {uri}")
            self.driver = GraphDatabase.driver(uri, auth=(user, password))
            self.driver.verify_connectivity()
            self.connected = True
            logger.info("成功连接到 Neo4j 数据库")
        except Exception as e:
            self.connected = False
            logger.error(f"Neo4j 连接失败: {str(e)}")
            self.driver = None

    def close(self):
        if self.driver:
            self.driver.close()
            logger.info("Neo4j 连接已关闭")

    # 创建节点和关系
    def create_entity(self, tx, entity_type, entity_id, name, properties=None):
        if properties is None:
            properties = {}
        
        query = f"""
            MERGE (e:{entity_type} {{id: $id}})
            ON CREATE SET e.name = $name
            """
        
        # 添加其他属性
        for key, value in properties.items():
            query += f", e.{key} = ${key}"
            
        params = {"id": entity_id, "name": name, **properties}
        tx.run(query, **params)

    def create_relationship(self, tx, source_id, target_id, rel_type, confidence, properties=None):
        if properties is None:
            properties = {}
            
        query = f"""
            MATCH (s {{id: $source_id}}), (t {{id: $target_id}})
            MERGE (s)-[r:{rel_type}]->(t)
            ON CREATE SET r.confidence = $confidence, r.created = datetime()
            """
            
        # 添加其他属性
        for key, value in properties.items():
            query += f", r.{key} = ${key}"
            
        params = {"source_id": source_id, "target_id": target_id, "confidence": confidence, **properties}
        tx.run(query, **params)

    # 获取图数据
    def get_graph_data(self, limit=500):
        if not self.connected:
            logger.warning("Neo4j 未连接，无法获取图数据")
            return {"nodes": [], "links": [], "error": "Neo4j database not connected"}
            
        try:
            with self.driver.session() as session:
                # 检查Neo4j中的节点总数
                count_result = session.run("MATCH (n) RETURN count(n) as node_count")
                node_count = count_result.single()["node_count"]
                logger.info(f"Neo4j数据库中发现 {node_count} 个节点")
                
                if node_count == 0:
                    logger.warning("Neo4j数据库中没有节点数据")
                    return {"nodes": [], "links": [], "error": "No nodes in database"}
                
                # 首先获取所有节点及其属性
                logger.info(f"开始获取节点数据，限制为 {limit} 个节点")
                node_result = session.run("""
                    MATCH (n)
                    RETURN id(n) as nodeId, n, labels(n) as types
                    LIMIT $limit
                """, limit=limit)
                
                nodes = []
                node_internal_ids_map = {}  # 映射内部ID到我们生成的ID
                
                for record in node_result:
                    node = record["n"]
                    node_internal_id = record["nodeId"]
                    node_type = record["types"][0] if record["types"] else "Unknown"
                    
                    # 生成一个一致的ID (使用内部ID字符串)
                    node_id = f"n{node_internal_id}"
                    
                    # 将节点内部ID映射到我们生成的ID
                    node_internal_ids_map[node_internal_id] = node_id
                    
                    # 转换节点属性到字典
                    node_dict = dict(node)
                    # 添加我们生成的ID和类型
                    node_dict["id"] = node_id
                    node_dict["type"] = node_type
                    # 添加节点名称(如果没有)
                    if "name" not in node_dict:
                        # 使用查找可能的名称属性
                        for name_attr in ["name", "title", "label", "value"]:
                            if name_attr in node_dict:
                                node_dict["name"] = node_dict[name_attr]
                                break
                        # 如果仍然没有名称，使用ID
                        if "name" not in node_dict:
                            node_dict["name"] = node_id
                    
                    nodes.append(node_dict)
                
                logger.info(f"成功获取 {len(nodes)} 个节点")
                
                # 然后获取关系，这次使用内部ID
                logger.info("开始获取关系数据")
                rel_result = session.run("""
                    MATCH (n)-[r]->(m)
                    RETURN id(n) AS source_internal_id, 
                           id(m) AS target_internal_id, 
                           type(r) AS type, 
                           properties(r) AS properties
                    LIMIT $limit
                """, limit=limit)
                
                links = []
                rel_count = 0
                for record in rel_result:
                    rel_count += 1
                    source_internal_id = record["source_internal_id"]
                    target_internal_id = record["target_internal_id"]
                    rel_type = record["type"]
                    
                    logger.info(f"找到关系: ({source_internal_id})-[{rel_type}]->({target_internal_id})")
                    
                    # 转换内部ID到我们的ID格式
                    if source_internal_id in node_internal_ids_map and target_internal_id in node_internal_ids_map:
                        link = {
                            "source": node_internal_ids_map[source_internal_id],
                            "target": node_internal_ids_map[target_internal_id],
                            "type": rel_type,
                            **record["properties"]
                        }
                        links.append(link)
                        logger.info(f"添加关系: {node_internal_ids_map[source_internal_id]} -> {node_internal_ids_map[target_internal_id]}, 类型: {rel_type}")
                    else:
                        logger.warning(f"找到一个关系，但其源节点或目标节点不在获取的节点列表中: {source_internal_id} -> {target_internal_id}")
                
                if rel_count == 0:
                    logger.warning("Neo4j数据库中没有找到任何关系")
                    # 尝试用其他查询方式获取关系
                    logger.info("尝试使用替代查询获取关系")
                    alt_rel_result = session.run("""
                        MATCH ()-[r]->() 
                        RETURN count(r) as count
                    """)
                    alt_rel_count = alt_rel_result.single()["count"]
                    logger.info(f"替代查询发现 {alt_rel_count} 个关系")
                    
                    if alt_rel_count > 0:
                        # 使用更简单的查询再次尝试
                        simple_rel_result = session.run("""
                            MATCH (n)-[r]->(m)
                            RETURN id(n) AS source_id, id(m) AS target_id, type(r) AS type
                            LIMIT $limit
                        """, limit=limit)
                        
                        simple_rel_count = 0
                        for record in simple_rel_result:
                            simple_rel_count += 1
                            logger.info(f"简单查询找到关系: {record['source_id']} -> {record['target_id']}, 类型: {record['type']}")
                        
                        logger.info(f"简单查询找到 {simple_rel_count} 个关系")
                
                logger.info(f"成功获取 {len(links)} 个关系")
                
                return {"nodes": nodes, "links": links}
        except Exception as e:
            logger.error(f"获取图数据失败: {str(e)}")
            import traceback
            logger.error(traceback.format_exc())
            return {"nodes": [], "links": [], "error": str(e)}

    # 评估三元组可信度
    def evaluate_trustworthiness(self, source_id, target_id, rel_type):
        """评估三元组(source, relation, target)的可信度"""
        try:
            confidence_score = mtmkg_model.evaluate_triple(source_id, target_id, rel_type)
            return confidence_score
        except Exception as e:
            logger.error(f"评估可信度失败: {str(e)}")
            return None

    # 获取所有知识图谱中的三元组数据
    def get_all_triples(self, limit=100):
        """获取知识图谱中所有的三元组数据"""
        if not self.connected:
            logger.warning("Neo4j 未连接，无法获取三元组数据")
            return []
            
        try:
            with self.driver.session() as session:
                result = session.run("""
                    MATCH (s)-[r]->(t)
                    WHERE NOT type(r) = 'MENTIONS'
                    RETURN s.id AS source_id, s.name AS source_name, labels(s)[0] AS source_type,
                           t.id AS target_id, t.name AS target_name, labels(t)[0] AS target_type,
                           type(r) AS relation_type, r.confidence AS confidence
                    ORDER BY confidence DESC
                    LIMIT $limit
                """, limit=limit)
                
                triples = []
                for record in result:
                    triple = {
                        "source_id": record["source_id"],
                        "source_name": record["source_name"],
                        "source_type": record["source_type"],
                        "target_id": record["target_id"],
                        "target_name": record["target_name"],
                        "target_type": record["target_type"],
                        "relation_type": record["relation_type"],
                        "confidence": record["confidence"] if record["confidence"] is not None else 0.5
                    }
                    triples.append(triple)
                    
                return triples
        except Exception as e:
            logger.error(f"获取三元组数据失败: {str(e)}")
            return []

    # 检查数据库中的节点和关系
    def check_database_entities(self):
        """检查Neo4j数据库中的节点和关系信息"""
        if not self.connected:
            logger.warning("Neo4j 未连接，无法检查数据库")
            return {"status": "error", "message": "Neo4j database not connected"}
            
        try:
            with self.driver.session() as session:
                # 检查节点数量
                node_count_result = session.run("MATCH (n) RETURN count(n) as count")
                node_count = node_count_result.single()["count"]
                
                # 检查节点属性
                node_props_result = session.run("""
                    MATCH (n) 
                    UNWIND keys(n) AS property
                    RETURN DISTINCT property, count(n) as count
                    ORDER BY count DESC
                """)
                node_properties = [{"property": record["property"], "count": record["count"]} 
                                 for record in node_props_result]
                
                # 检查节点标签
                node_labels_result = session.run("""
                    MATCH (n)
                    RETURN DISTINCT labels(n) as labels, count(*) as count
                    ORDER BY count DESC
                """)
                node_labels = []
                for record in node_labels_result:
                    labels = record["labels"]
                    count = record["count"]
                    node_labels.append({"labels": labels, "count": count})
                
                # 检查关系数量
                rel_count_result = session.run("MATCH ()-[r]->() RETURN count(r) as count")
                rel_count = rel_count_result.single()["count"]
                
                # 检查关系类型
                rel_types_result = session.run("""
                    MATCH ()-[r]->()
                    RETURN DISTINCT type(r) as type, count(r) as count
                    ORDER BY count DESC
                """)
                rel_types = [{"type": record["type"], "count": record["count"]} 
                           for record in rel_types_result]
                
                # 检查示例节点
                sample_nodes_result = session.run("""
                    MATCH (n)
                    RETURN id(n) as id, n, labels(n) as labels
                    LIMIT 5
                """)
                sample_nodes = []
                for record in sample_nodes_result:
                    node = record["n"]
                    node_id = record["id"]
                    labels = record["labels"]
                    sample_nodes.append({
                        "internal_id": node_id,
                        "labels": labels,
                        "properties": dict(node)
                    })
                
                # 检查示例关系
                sample_rels_result = session.run("""
                    MATCH (n)-[r]->(m)
                    RETURN id(n) as source_id, labels(n) as source_labels,
                           id(m) as target_id, labels(m) as target_labels,
                           type(r) as rel_type, properties(r) as properties
                    LIMIT 5
                """)
                sample_relationships = []
                for record in sample_rels_result:
                    sample_relationships.append({
                        "source_id": record["source_id"],
                        "source_labels": record["source_labels"],
                        "target_id": record["target_id"],
                        "target_labels": record["target_labels"],
                        "type": record["rel_type"],
                        "properties": record["properties"]
                    })
                
                return {
                    "status": "success",
                    "node_count": node_count,
                    "relationship_count": rel_count,
                    "node_properties": node_properties,
                    "node_labels": node_labels,
                    "relationship_types": rel_types,
                    "sample_nodes": sample_nodes,
                    "sample_relationships": sample_relationships
                }
        except Exception as e:
            logger.error(f"检查数据库实体失败: {str(e)}")
            import traceback
            logger.error(traceback.format_exc())
            return {
                "status": "error",
                "message": f"检查数据库实体失败: {str(e)}"
            }

    # 专门为12k_all.txt格式的数据获取图谱
    def get_entity_relations(self, limit=5000):
        """获取实体间的关系，特别针对Entity类型节点之间的关系"""
        if not self.connected:
            logger.warning("Neo4j 未连接，无法获取实体关系")
            return {"nodes": [], "links": [], "error": "Neo4j database not connected"}
            
        try:
            with self.driver.session() as session:
                # 检查实体节点数量
                count_result = session.run("""
                    MATCH (n) 
                    RETURN count(n) as count, 
                           count(DISTINCT labels(n)) as label_count
                """)
                record = count_result.single()
                node_count = record["count"]
                label_count = record["label_count"]
                logger.info(f"Neo4j数据库中发现 {node_count} 个节点，{label_count} 种标签类型")
                
                # 检查是否有关系
                rel_count_result = session.run("MATCH ()-[r]->() RETURN count(r) as count")
                rel_count = rel_count_result.single()["count"]
                logger.info(f"Neo4j数据库中发现 {rel_count} 个关系")
                
                # 获取所有节点，使用更简单的查询
                logger.info(f"开始获取全部节点")
                nodes_result = session.run("""
                    MATCH (n)
                    RETURN id(n) as id, n, labels(n) as types
                    LIMIT $limit
                """, limit=limit)
                
                nodes = []
                node_id_map = {}  # 内部ID到生成ID的映射
                
                for record in nodes_result:
                    node = record["n"]
                    node_id = record["id"]
                    node_types = record["types"]
                    
                    # 生成唯一ID
                    generated_id = f"n{node_id}"
                    node_id_map[node_id] = generated_id
                    
                    # 构建节点对象
                    properties = dict(node)
                    # 序列化可能存在的Neo4j特殊类型
                    properties = serialize_neo4j_data(properties)
                    node_type = node_types[0] if node_types else "Unknown"
                    
                    # 设置节点类型
                    node_obj = {
                        "id": generated_id,
                        "type": node_type,
                        **properties
                    }
                    
                    # 确保有名称
                    if "name" not in node_obj:
                        # 尝试从其他属性中找名称
                        name_found = False
                        for key, value in properties.items():
                            if isinstance(value, str) and len(value) < 50:
                                node_obj["name"] = value
                                name_found = True
                                break
                        
                        if not name_found:
                            node_obj["name"] = generated_id
                    
                    nodes.append(node_obj)
                
                logger.info(f"成功获取 {len(nodes)} 个节点")
                
                # 获取所有关系，不应用任何过滤器
                logger.info("开始获取全部关系")
                rels_result = session.run("""
                    MATCH (n)-[r]->(m)
                    RETURN id(n) as source_id, id(m) as target_id, 
                           type(r) as rel_type, properties(r) as props
                    LIMIT $limit
                """, limit=limit)
                
                links = []
                added_relations = set()  # 用于跟踪已添加的关系，避免重复
                
                for record in rels_result:
                    source_id = record["source_id"]
                    target_id = record["target_id"]
                    rel_type = record["rel_type"]
                    props = record["props"]
                    
                    # 序列化Neo4j特殊类型(如DateTime)
                    props = serialize_neo4j_data(props)
                    
                    if source_id in node_id_map and target_id in node_id_map:
                        source = node_id_map[source_id]
                        target = node_id_map[target_id]
                        
                        # 避免重复关系
                        relation_key = f"{source}:{rel_type}:{target}"
                        if relation_key not in added_relations:
                            links.append({
                                "source": source,
                                "target": target,
                                "type": rel_type,
                                **props
                            })
                            added_relations.add(relation_key)
                            logger.info(f"添加关系: {source} -[{rel_type}]-> {target}")
                    else:
                        logger.warning(f"关系的端点不在节点集中: {source_id} -> {target_id}")
                
                logger.info(f"成功获取 {len(links)} 个关系")
                
                if len(links) == 0 and rel_count > 0:
                    logger.warning("检测到数据库有关系但查询返回0个关系，尝试创建基本关系...")
                    
                    # 如果没有关系但数据库中有节点，创建一些基本关系用于测试
                    if len(nodes) > 1:
                        for i in range(min(20, len(nodes)-1)):
                            source = nodes[i]["id"]
                            target = nodes[i+1]["id"]
                            links.append({
                                "source": source,
                                "target": target,
                                "type": "RELATED_TO"
                            })
                        logger.info(f"创建了 {min(20, len(nodes)-1)} 个测试关系")
                
                return {"nodes": nodes, "links": links}
        except Exception as e:
            logger.error(f"获取实体关系失败: {str(e)}")
            import traceback
            logger.error(traceback.format_exc())
            return {"nodes": [], "links": [], "error": str(e)}

    # 获取所有实体类型
    def get_entity_types(self):
        """获取知识图谱中所有的实体类型"""
        if not self.connected:
            logger.warning("Neo4j 未连接，无法获取实体类型")
            return []
            
        try:
            with self.driver.session() as session:
                result = session.run("""
                    MATCH (n)
                    RETURN DISTINCT labels(n)[0] as type, count(*) as count
                    ORDER BY count DESC
                """)
                
                types = []
                for record in result:
                    entity_type = record["type"]
                    if entity_type:
                        types.append({
                            "type": entity_type,
                            "count": record["count"]
                        })
                
                return types
        except Exception as e:
            logger.error(f"获取实体类型失败: {str(e)}")
            return []

# 创建 Knowledge Graph 实例
graph = KnowledgeGraph(URI, USER, PASSWORD)

models.Base.metadata.create_all(bind=engine)

# 定义应用生命周期
@asynccontextmanager
async def lifespan(app: FastAPI):
    # 应用启动时的操作
    logger.info("应用启动，进行初始化...")
    
    # 测试Neo4j连接
    if graph.connected:
        logger.info("Neo4j连接正常")
    else:
        logger.error("Neo4j连接失败")
    
    yield  # 应用运行
    
    # 应用关闭时的操作
    logger.info("应用关闭，清理资源...")

    # 关闭Neo4j连接
    graph.close()

# 创建 FastAPI 应用实例
app = FastAPI(lifespan=lifespan)

# 添加 CORS 中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 在生产环境中，应该限制为前端的域名
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 基本的健康检查端点
@app.get("/health")
def health_check():
    return {
        "status": "ok",
        "neo4j_connected": graph.connected,
        "timestamp": datetime.now().isoformat()
    }

# Neo4j数据库检查端点
@app.get("/check-neo4j")
def check_neo4j():
    """检查Neo4j数据库状态和数据"""
    if not graph.connected:
        return {
            "status": "error",
            "message": "Neo4j数据库未连接",
            "neo4j_uri": URI
        }
        
    try:
        with graph.driver.session() as session:
            # 检查节点总数
            node_count = session.run("MATCH (n) RETURN count(n) as count").single()["count"]
            
            # 获取节点类型统计
            type_result = session.run("""
                MATCH (n) 
                WITH labels(n)[0] as type, count(*) as count 
                RETURN type, count 
                ORDER BY count DESC
            """)
            node_types = [{"type": record["type"], "count": record["count"]} for record in type_result]
            
            # 获取关系类型统计
            rel_result = session.run("""
                MATCH ()-[r]->() 
                WITH type(r) as type, count(*) as count 
                RETURN type, count 
                ORDER BY count DESC
            """)
            rel_types = [{"type": record["type"], "count": record["count"]} for record in rel_result]
            
            # 检查示例节点
            sample_nodes = []
            if node_count > 0:
                sample_result = session.run("""
                    MATCH (n) 
                    RETURN n, labels(n) as types 
                    LIMIT 5
                """)
                for record in sample_result:
                    node = record["n"]
                    node_type = record["types"][0] if record["types"] else "Unknown"
                    sample_nodes.append({
                        "id": node.get("id", str(node.id)),
                        "type": node_type,
                        "properties": dict(node)
                    })
            
            return {
                "status": "ok",
                "neo4j_uri": URI,
                "node_count": node_count,
                "node_types": node_types,
                "relationship_types": rel_types,
                "sample_nodes": sample_nodes
            }
    except Exception as e:
        logger.error(f"检查Neo4j数据库时出错: {str(e)}")
        import traceback
        logger.error(traceback.format_exc())
        return {
            "status": "error",
            "message": f"检查Neo4j数据库时出错: {str(e)}",
            "neo4j_uri": URI
        }

# 检查Neo4j实体详情端点
@app.get("/check-entities")
def check_database_entities(current_user: models.User = Depends(auth.get_current_user)):
    """检查Neo4j数据库中的实体详细信息"""
    result = graph.check_database_entities()
    return result

# 用户注册
@app.post("/register", response_model=schemas.User)
def register(user: schemas.UserCreate, db: Session = Depends(get_db)):
    # 检查用户是否已存在
    db_user = db.query(models.User).filter(models.User.username == user.username).first()
    if db_user:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail="用户名已存在"
        )
    
    # 创建新用户
    hashed_password = auth.get_password_hash(user.password)
    db_user = models.User(username=user.username, email=user.email, password_hash=hashed_password)
    db.add(db_user)
    db.commit()
    db.refresh(db_user)
    return db_user

# 用户登录
@app.post("/token", response_model=schemas.Token)
def login(user_credentials: schemas.UserLogin, db: Session = Depends(get_db)):
    user = auth.authenticate_user(db, user_credentials.username, user_credentials.password)
    if not user:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="用户名或密码错误",
            headers={"WWW-Authenticate": "Bearer"},
        )
    
    access_token = auth.create_access_token(data={"sub": user.username})
    return {"access_token": access_token, "token_type": "bearer"}

# 获取当前用户信息
@app.get("/users/me", response_model=schemas.User)
def read_users_me(current_user: models.User = Depends(auth.get_current_user)):
    return current_user

# 获取图谱数据
@app.get("/graph", dependencies=[Depends(auth.get_current_user)])
def get_graph(search: str = "", node_type: str = "", relation_types: str = "", limit: int = 100):
    """获取知识图谱数据"""
    try:
        with graph.driver.session() as session:
            # 构建查询条件
            conditions = []
            parameters = {}
            
            if search:
                conditions.append("n.name CONTAINS $search OR m.name CONTAINS $search")
                parameters["search"] = search
            
            if node_type:
                node_types = node_type.split(',')
                node_type_conditions = []
                for i, nt in enumerate(node_types):
                    param_name = f"node_type_{i}"
                    node_type_conditions.append(f"n:{nt} OR m:{nt}")
                    parameters[param_name] = nt
                if node_type_conditions:
                    conditions.append("(" + " OR ".join(node_type_conditions) + ")")
            
            relation_filter = ""
            if relation_types:
                rel_types = relation_types.split(',')
                rel_conditions = []
                for i, rt in enumerate(rel_types):
                    param_name = f"rel_type_{i}"
                    rel_conditions.append(f"type(r) = ${param_name}")
                    parameters[param_name] = rt
                if rel_conditions:
                    relation_filter = "WHERE " + " OR ".join(rel_conditions)
            
            where_clause = ""
            if conditions:
                where_clause = "WHERE " + " AND ".join(conditions)
            
            # 构建完整查询
            query = f"""
                MATCH (n)-[r]->(m)
                {relation_filter}
                {where_clause}
                RETURN n, r, m, type(r) as relation_type
                LIMIT {limit}
            """
            
            result = session.run(query, parameters)
            
            nodes = {}
            links = []
            relation_types_set = set()
            
            for record in result:
                source_node = record["n"]
                target_node = record["m"]
                relationship = record["r"]
                relation_type = record["relation_type"]
                
                # 记录关系类型
                relation_types_set.add(relation_type)
                
                # 处理源节点
                source_id = str(source_node.id)
                if source_id not in nodes:
                    source_props = dict(source_node)
                    source_props["id"] = f"n{source_id}"
                    source_props["type"] = list(source_node.labels)[0] if source_node.labels else "Unknown"
                    source_props["name"] = source_props.get("name", f"Node-{source_id}")
                    nodes[source_id] = source_props
                
                # 处理目标节点
                target_id = str(target_node.id)
                if target_id not in nodes:
                    target_props = dict(target_node)
                    target_props["id"] = f"n{target_id}"
                    target_props["type"] = list(target_node.labels)[0] if target_node.labels else "Unknown"
                    target_props["name"] = target_props.get("name", f"Node-{target_id}")
                    nodes[target_id] = target_props
                
                # 处理关系，确保同时包含type和original_relation
                relationship_props = dict(relationship)
                links.append({
                    "source": f"n{source_id}",
                    "target": f"n{target_id}",
                    "type": relationship.type,
                    "original_relation": relation_type,  # 添加原始关系类型
                    **relationship_props
                })
            
            # 记录查询到的关系类型
            logger.info(f"查询条件下发现的关系类型: {relation_types_set}")
            
            return {
                "nodes": list(nodes.values()),
                "links": links,
                "relationship_types": list(relation_types_set),
                "message": f"Found {len(nodes)} nodes and {len(links)} relationships"
            }
    except Exception as e:
        logger.error(f"获取图谱数据失败: {str(e)}")
        import traceback
        return {
            "nodes": [],
            "links": [],
            "error": str(e),
            "traceback": traceback.format_exc()
        }

# 获取实体关系
@app.get("/entity-graph")
def get_entity_graph(current_user: models.User = Depends(auth.get_current_user)) -> Dict:
    """获取特别优化用于显示实体关系的图谱数据"""
    graph_data = graph.get_entity_relations()
    return graph_data

# 评估三元组可信度
@app.post("/evaluate")
def evaluate_trust(
    eval_input: TrustEvaluationInput,
    current_user: models.User = Depends(auth.get_current_user)
) -> Dict:
    confidence = graph.evaluate_trustworthiness(
        eval_input.source_id,
        eval_input.target_id,
        eval_input.relation_type
    )
    
    return {"confidence": confidence}

# 获取所有三元组数据
@app.get("/triples")
def get_triples(
    limit: int = 1000, 
    current_user: models.User = Depends(auth.get_current_user)
) -> List:
    return graph.get_all_triples(limit)

# 评估所有三元组可信度
@app.post("/evaluate/all")
def evaluate_all_triples(
    current_user: models.User = Depends(auth.get_current_user)
) -> Dict:
    """评估知识图谱中所有三元组的可信度"""
    try:
        # 获取所有三元组
        triples = graph.get_all_triples(limit=1000)
        
        # 评估每个三元组的可信度
        updated_count = 0
        for triple in triples:
            confidence = mtmkg_model.evaluate_triple(
                triple["source_id"], 
                triple["target_id"], 
                triple["relation_type"]
            )
            
            if confidence is not None:
                updated_count += 1
                
            return {
            "status": "success",
            "message": f"成功评估 {updated_count}/{len(triples)} 个三元组的可信度",
            "updated_count": updated_count,
            "total_count": len(triples)
        }
        
    except Exception as e:
        logger.error(f"批量评估可信度失败: {str(e)}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"评估可信度失败: {str(e)}"
        )

# 不需要认证的API端点，用于测试
@app.get("/test-graph")
def test_graph():
    """测试用的图谱数据端点，无需认证"""
    try:
        with graph.driver.session() as session:
            # 修改查询，获取所有与Dragonfly相关的节点和关系，包括入边和出边
            # 首先获取Dragonfly节点
            find_node_query = """
                MATCH (n:Threat_actor)
                WHERE n.name='Dragonfly' OR n.alias='Dragonfly'
                RETURN n
            """
            
            node_result = session.run(find_node_query)
            dragonfly_node = None
            for record in node_result:
                dragonfly_node = record["n"]
                break
                
            if not dragonfly_node:
                # 尝试寻找任何标签的Dragonfly节点
                wider_query = """
                    MATCH (n)
                    WHERE n.name='Dragonfly' OR n.alias='Dragonfly'
                    RETURN n
                """
                node_result = session.run(wider_query)
                for record in node_result:
                    dragonfly_node = record["n"]
                    break
            
            if not dragonfly_node:
                logger.warning("未找到名为Dragonfly的节点")
                return {
                    "nodes": [],
                    "links": [],
                    "error": "未找到名为Dragonfly的节点"
                }
                
            dragonfly_id = dragonfly_node.id
            logger.info(f"找到Dragonfly节点，内部ID: {dragonfly_id}")
            
            # 获取所有与Dragonfly相关的关系 - 包括入边和出边
            relations_query = """
                MATCH (n)-[r]-(m)
                WHERE id(n) = $node_id
                RETURN n, r, m, type(r) as relation_type, startNode(r) as source, endNode(r) as target
            """
            relations_result = session.run(relations_query, node_id=dragonfly_id)
            
            nodes = {}
            links = []
            relation_types = set()
            related_node_ids = set()  # 跟踪所有关联节点的ID
            
            # 添加Dragonfly节点
            dragonfly_props = dict(dragonfly_node)
            dragonfly_props["id"] = f"n{dragonfly_id}"
            dragonfly_props["type"] = list(dragonfly_node.labels)[0] if dragonfly_node.labels else "Unknown"
            dragonfly_props["name"] = dragonfly_props.get("name", "Dragonfly")
            nodes[str(dragonfly_id)] = dragonfly_props
            related_node_ids.add(dragonfly_id)  # 添加Dragonfly到关联节点集合
            
            for record in relations_result:
                source_node = record["source"]
                target_node = record["target"]
                relationship = record["r"]
                relation_type = record["relation_type"]
                
                # 记录关系类型
                relation_types.add(relation_type)
                
                # 处理源节点
                source_id = source_node.id
                if str(source_id) not in nodes:
                    source_props = dict(source_node)
                    source_props["id"] = f"n{source_id}"
                    source_props["type"] = list(source_node.labels)[0] if source_node.labels else "Unknown"
                    source_props["name"] = source_props.get("name", f"Node-{source_id}")
                    nodes[str(source_id)] = source_props
                related_node_ids.add(source_id)  # 添加到关联节点集合
                
                # 处理目标节点
                target_id = target_node.id
                if str(target_id) not in nodes:
                    target_props = dict(target_node)
                    target_props["id"] = f"n{target_id}"
                    target_props["type"] = list(target_node.labels)[0] if target_node.labels else "Unknown"
                    target_props["name"] = target_props.get("name", f"Node-{target_id}")
                    nodes[str(target_id)] = target_props
                related_node_ids.add(target_id)  # 添加到关联节点集合
                
                # 处理关系，确保同时包含type和original_relation
                relationship_props = dict(relationship)
                link_key = f"{source_id}_{relationship.type}_{target_id}"
                links.append({
                    "source": f"n{source_id}",
                    "target": f"n{target_id}",
                    "type": relationship.type,
                    "original_relation": relation_type,  # 添加原始关系类型，确保前端可以访问
                    "key": link_key,  # 添加唯一键，避免重复关系
                    **relationship_props
                })
            
            # 第二步：获取这些关联节点之间的关系（二度关系）
            logger.info(f"找到 {len(related_node_ids)} 个与Dragonfly关联的节点，现在查询它们之间的关系")
            
            # 将相关节点ID转换为列表
            related_ids_list = list(related_node_ids)
            
            # 由于Neo4j参数不能直接使用列表，我们需要分批处理
            batch_size = 50  # 每批处理的节点数
            added_links = set()  # 跟踪已添加的关系，避免重复
            
            for i in range(0, len(related_ids_list), batch_size):
                batch = related_ids_list[i:i+batch_size]
                
                # 对于当前批次中的每个节点，查询其与其他关联节点的关系
                for node_id in batch:
                    secondary_query = """
                        MATCH (n)-[r]-(m)
                        WHERE id(n) = $source_id AND id(m) IN $target_ids
                        RETURN n, r, m, type(r) as relation_type, startNode(r) as source, endNode(r) as target
                    """
                    # 排除当前节点，只查询其与其他关联节点的关系
                    target_ids = [tid for tid in related_ids_list if tid != node_id]
                    
                    if not target_ids:
                        continue  # 如果没有其他关联节点，跳过
                    
                    secondary_result = session.run(secondary_query, source_id=node_id, target_ids=target_ids)
                    
                    for record in secondary_result:
                        source_node = record["source"]
                        target_node = record["target"]
                        relationship = record["r"]
                        relation_type = record["relation_type"]
                        
                        # 记录关系类型
                        relation_types.add(relation_type)
                        
                        source_id = source_node.id
                        target_id = target_node.id
                        
                        # 创建关系的唯一键
                        link_key = f"{source_id}_{relationship.type}_{target_id}"
                        
                        # 检查这个关系是否已经添加过
                        if link_key in added_links:
                            continue
                            
                        # 记录这个关系已被添加
                        added_links.add(link_key)
                        
                        # 添加关系
                        relationship_props = dict(relationship)
                        links.append({
                            "source": f"n{source_id}",
                            "target": f"n{target_id}",
                            "type": relationship.type,
                            "original_relation": relation_type,
                            "key": link_key,
                            **relationship_props
                        })
            
            # 移除链接中的临时key字段
            for link in links:
                if 'key' in link:
                    del link['key']
            
            # 记录所有关系类型，用于调试
            logger.info(f"发现关系类型: {relation_types}")
            logger.info(f"节点数量: {len(nodes.values())}, 关系数量: {len(links)}")
            
            return {
                "nodes": list(nodes.values()),
                "links": links,
                "relationship_types": list(relation_types),
                "message": f"Found {len(nodes)} nodes and {len(links)} relationships with {len(relation_types)} relation types"
            }
    except Exception as e:
        import traceback
        return {
            "nodes": [],
            "links": [],
            "error": str(e),
            "traceback": traceback.format_exc()
        }

# 获取所有实体类型
@app.get("/entity-types")
def get_entity_types():
    """获取所有实体类型，不需要认证"""
    try:
        types = graph.get_entity_types()
        return types
    except Exception as e:
        logger.error(f"获取实体类型失败: {str(e)}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"获取实体类型失败: {str(e)}"
        )

# 通过名称查询单个节点及其关系
@app.get("/node-by-name/{name}")
def get_node_by_name(name: str):
    """根据节点名称查询节点及其直接关联的关系"""
    try:
        with graph.driver.session() as session:
            # 查询具有指定名称的节点 - 扩展查询以匹配更多可能的属性
            node_query = """
                MATCH (n) 
                WHERE n.name = $name OR n.alias = $name
                RETURN n
            """
            node_result = session.run(node_query, name=name)
            
            # 检查是否找到节点
            node_record = node_result.single()
            if not node_record:
                return {
                    "error": f"未找到名称为 '{name}' 的节点",
                    "nodes": [],
                    "links": []
                }
            
            node = node_record["n"]
            node_id = node.id
            node_props = dict(node)
            node_props["id"] = f"n{node_id}"
            node_props["type"] = list(node.labels)[0] if node.labels else "Unknown"
            
            # 查询与此节点相关的所有关系 - 使用id匹配以确保准确性
            relations_query = """
                MATCH (n)-[r]-(m)
                WHERE id(n) = $node_id
                RETURN n, r, m, type(r) as relation_type, startNode(r) as source, endNode(r) as target
            """
            relations_result = session.run(relations_query, node_id=node_id)
            
            # 处理节点和关系
            nodes = {str(node_id): node_props}
            links = []
            relation_types = set()
            
            for record in relations_result:
                source_node = record["source"]
                target_node = record["target"]
                relationship = record["r"]
                relation_type = record["relation_type"]
                
                # 记录关系类型
                relation_types.add(relation_type)
                
                # 处理相关节点
                source_id = str(source_node.id)
                if source_id not in nodes:
                    source_props = dict(source_node)
                    source_props["id"] = f"n{source_id}"
                    source_props["type"] = list(source_node.labels)[0] if source_node.labels else "Unknown"
                    source_props["name"] = source_props.get("name", f"Node-{source_id}")
                    nodes[source_id] = source_props
                
                target_id = str(target_node.id)
                if target_id not in nodes:
                    target_props = dict(target_node)
                    target_props["id"] = f"n{target_id}"
                    target_props["type"] = list(target_node.labels)[0] if target_node.labels else "Unknown"
                    target_props["name"] = target_props.get("name", f"Node-{target_id}")
                    nodes[target_id] = target_props
                
                # 处理关系
                relationship_props = dict(relationship)
                
                links.append({
                    "source": f"n{source_id}",
                    "target": f"n{target_id}",
                    "type": relationship.type,
                    "original_relation": relation_type,
                    **relationship_props
                })
            
            logger.info(f"查询节点 '{name}' 的结果: 找到 {len(nodes)} 个节点和 {len(links)} 个关系")
            logger.info(f"关系类型: {relation_types}")
            
            return {
                "nodes": list(nodes.values()),
                "links": links,
                "relationship_types": list(relation_types),
                "message": f"Found node '{name}' with {len(links)} relationships"
            }
    except Exception as e:
        logger.error(f"查询节点 '{name}' 失败: {str(e)}")
        import traceback
        logger.error(traceback.format_exc())
        return {
            "error": str(e),
            "nodes": [],
            "links": []
        }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)